import torch
import matplotlib.pyplot as plt

# 设置中文字体，解决显示中文乱码问题
import matplotlib

matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 使用 SimHei 字体，或者尝试 Microsoft YaHei
matplotlib.rcParams['axes.unicode_minus'] = False  # 防止负号显示问题

def visualize_featmap(featmap_before, featmap_after, idx=0, cmap="viridis"):
    """
    可视化特征图在应用旋转矩阵前后的变化。

    :param featmap_before: 应用旋转矩阵前的特征图 (torch.Tensor)
    :param featmap_after: 应用旋转矩阵后的特征图 (torch.Tensor)
    :param idx: 选择可视化的通道索引 (int)
    :param cmap: 热力图的颜色映射 (str)
    """
    # 如果特征图是 4D，选择第一个 batch 和一个通道
    if len(featmap_before.shape) == 4:
        featmap_before = featmap_before[0, idx].detach().cpu().numpy()
        featmap_after = featmap_after[0, idx].detach().cpu().numpy()
    elif len(featmap_before.shape) == 3:  # 如果是 3D 张量
        featmap_before = featmap_before[idx].detach().cpu().numpy()
        featmap_after = featmap_after[idx].detach().cpu().numpy()
    else:
        raise ValueError("特征图的维度应为 3D 或 4D")

    # 绘制特征图
    plt.figure(figsize=(12, 6))

    # 原始特征图
    plt.subplot(1, 2, 1)
    plt.imshow(featmap_before, cmap=cmap)
    plt.title("特征图 - 旋转矩阵前")
    plt.colorbar()

    # 应用旋转矩阵后的特征图
    plt.subplot(1, 2, 2)
    plt.imshow(featmap_after, cmap=cmap)
    plt.title("特征图 - 旋转矩阵后")
    plt.colorbar()

    plt.tight_layout()
    plt.show()


# 示例：加载特征图和旋转矩阵
featmap = torch.rand((1, 3, 64, 64))  # 示例特征图 (batch_size, channels, height, width)
rotation_matrix = torch.eye(3)  # 示例旋转矩阵 (3x3)

# 计算旋转矩阵作用后的特征图
featmap_flattened = featmap.view(1, 3, -1)  # 展平为 (batch_size, channels, height*width)
featmap_transformed = torch.matmul(rotation_matrix, featmap_flattened).view(1, 3, 64, 64)  # 重新reshape为原形状

# 可视化
visualize_featmap(featmap, featmap_transformed, idx=0, cmap="viridis")
